import sys
sys.path.append("core_data_generation")
import numpy as np
import matplotlib.pyplot as plt
from model_var_disc_vectorized import *
from Parameters_1D import spatialVariables_1D_act, temporalVariable
from model_simulation import simulate_Richards_equation
import time
import copy 
init_Time = time.process_time()
soilPars = pars_mz1()
totalDepth, axialNodes, totalNodes= spatialVariables_1D_act()
samplingTime, samplingTimeInternal, internalTimeSteps, totTimeSteps = temporalVariable() #time parameters

def compute_root_zone_moisture(Xk, rooting_depth):
    if rooting_depth == 1.00:    
        root_zone_head = 0.10*((1/6)*(Xk[0] + Xk[1] + Xk[2] + Xk[3] + Xk[4] + Xk[5]))  + \
                         0.20*((1/6)*(Xk[5] + Xk[6] + Xk[7] + Xk[8] + Xk[9] + Xk[10])) + \
                         0.30*((1/11)*(Xk[10] + Xk[11] + Xk[12] + Xk[13] + Xk[14] + Xk[15] + Xk[16] + Xk[17] + Xk[18] + Xk[19] + Xk[20])) + \
                         0.40*((1/11)*(Xk[20] + Xk[21] + Xk[22] + Xk[23] + Xk[24] + Xk[25] + Xk[26] + Xk[27] + Xk[28] + Xk[29] + Xk[30]))

                         
    else :
        root_zone_head = 0.10*((1/6)*(Xk[10] + Xk[11] + Xk[12] + Xk[13] + Xk[14] + Xk[15])) +\
                         0.20*((1/6)*(Xk[15] + Xk[16] + Xk[17] + Xk[18] + Xk[19] + Xk[20])) +\
                         0.30*((1/6)*(Xk[20] + Xk[21] + Xk[22] + Xk[23] + Xk[24] + Xk[25])) + \
                         0.40*((1/6)*(Xk[25] + Xk[26] + Xk[27] + Xk[28] + Xk[29] + Xk[30]))
        
    return root_zone_head                                                                                                                                                                             

def obtain_soil_moisture(head):
    vol_moist = volMoistureAllNodes(head, soilPars)
    return vol_moist

def ObtainSimulationMetrics():
    sim_period = np.random.randint(low=3, high=8)        
    if sim_period == 3:
        irrig_rate = np.random.uniform(-6.0e-07, -4.5e-08)
        ref_evap = np.random.uniform(1.2e-08, 1.04e-07, sim_period)
        crop_coeff = np.random.uniform(0.4, 1.02, sim_period)
        rooting_depth = np.random.choice([0.50, 1.00])*np.ones(sim_period)
        return sim_period, irrig_rate, ref_evap, crop_coeff, rooting_depth
    
    elif sim_period == 4:
        irrig_rate = np.random.uniform(-6.0e-07, -4.5e-08)
        ref_evap = np.random.uniform(1.2e-08, 1.04e-07, sim_period)
        crop_coeff = np.random.uniform(0.4, 1.02, sim_period)
        rooting_depth = np.random.choice([0.50,  1.00])*np.ones(sim_period)
        return sim_period, irrig_rate, ref_evap, crop_coeff, rooting_depth
    
    elif sim_period == 5:
        irrig_rate = np.random.uniform(-6.0e-07, -4.5e-08)
        ref_evap = np.random.uniform(1.2e-08, 1.04e-07, sim_period)
        crop_coeff = np.random.uniform(0.4, 1.02, sim_period)
        rooting_depth = np.random.choice([0.50, 1.00])*np.ones(sim_period)
        return sim_period, irrig_rate, ref_evap, crop_coeff, rooting_depth
    
    elif sim_period == 6:
        irrig_rate = np.random.uniform(-6.0e-07, -4.5e-08)
        ref_evap = np.random.uniform(1.2e-08, 1.04e-07, sim_period)
        crop_coeff = np.random.uniform(0.4, 1.02, sim_period)
        rooting_depth = np.random.choice([0.50,1.00])*np.ones(sim_period)
        return sim_period, irrig_rate, ref_evap, crop_coeff, rooting_depth
    
    else:
        irrig_rate = np.random.uniform(-6.0e-07, -4.5e-08)
        ref_evap = np.random.uniform(1.2e-08, 1.04e-07, sim_period)
        crop_coeff = np.random.uniform(0.4, 1.02, sim_period)
        rooting_depth = np.random.choice([0.50, 1.00])*np.ones(sim_period)
        return sim_period, irrig_rate, ref_evap, crop_coeff, rooting_depth
    
#Run this function for a number of times and store the results
irrig_rate_train = []
ref_evap_train = []
crop_coeff_train = []
rooting_depth_train = []
sim_period_train = []

for i in range(28000):
    print(i)
    sim_period, irrig_rate, ref_evap, crop_coeff, rooting_depth= ObtainSimulationMetrics()
    if i in range(0, 1100):
        rooting_depth = [0.50]*sim_period
    elif i in range(1100, 2200):
        rooting_depth = [1.00]*sim_period        
    elif i in range(2200, 3300):
        rooting_depth = [0.50]*sim_period
    elif i in range(3300, 4400):
        rooting_depth = [1.00]*sim_period
        
    elif i in range(4400, 5500):
        rooting_depth = [0.50]*sim_period    
    elif i in range(5500, 6600):
        rooting_depth = [1.00]*sim_period
    elif i in range(6600, 7700):
        rooting_depth = [0.50]*sim_period
    elif i in range(7700, 8800):
        rooting_depth = [1.00]*sim_period
            
    elif i in range(8800, 9900):
        rooting_depth = [0.50]*sim_period
    elif i in range(9900, 11000):
        rooting_depth = [1.00]*sim_period
    elif i in range(11000, 12100):
        rooting_depth = [0.50]*sim_period
    elif i in range(12100, 13200):
        rooting_depth = [1.00]*sim_period
        
    elif i in range(13200, 14300):
        rooting_depth = [0.50]*sim_period
    elif i in range(14300, 15400):
        rooting_depth = [1.00]*sim_period
    elif i in range(15400, 16500):
        rooting_depth = [0.50]*sim_period
    elif i in range(16500, 17600):
        rooting_depth = [1.00]*sim_period
    
    elif i in range(17600, 18700):
        rooting_depth = [0.50]*sim_period
    elif i in range(18700, 19800):
        rooting_depth = [1.00]*sim_period
    elif i in range(19800, 20900):
        rooting_depth = [0.50]*sim_period
    elif i in range(20900, 22000):
        rooting_depth = [1.00]*sim_period
        
    elif i in range(22000, 23100):
        rooting_depth = [0.50]*sim_period
    elif i in range(23100, 24200):
        rooting_depth = [1.00]*sim_period
    elif i in range(24200, 25300):
        rooting_depth = [0.50]*sim_period
    else:
        rooting_depth = [1.00]*sim_period
        
    irrig_rate_ = [0]*sim_period
    irrig_rate_[0] = irrig_rate
    irrig_rate_train.extend(irrig_rate_)
    ref_evap_train.extend(ref_evap)
    crop_coeff_train.extend(crop_coeff)
    rooting_depth_train.extend(rooting_depth)
    sim_period_train.append(sim_period)    
    pass

#Perform the openloop simulations
#Define the initial condition 
x0_init = -0.96*np.ones(totalNodes)
data_set_all = np.zeros((len(crop_coeff_train), 6))
data_set_all_head = np.zeros((len(crop_coeff_train), 6))

_irrig_rates = np.array(irrig_rate_train)
_ref_evaps = np.array(ref_evap_train)
_crop_coeffs = np.array(crop_coeff_train)
_rooting_depth = np.array(rooting_depth_train)
_crop_coeffs = np.loadtxt('./core_data_generation/crop_coeff_extended.txt')
_lai_factors = np.loadtxt("./core_data_generation/LAI_factors.txt")

for i in range(len(_crop_coeffs)):
    if _crop_coeffs[i] < 0.20:
        _crop_coeffs[i] = 0.20
        pass
    pass
     
current_head = x0_init
for i in range(len(_irrig_rates)):
    print(i)   
    next_head =  simulate_Richards_equation(current_head, _irrig_rates[i], _ref_evaps[i], _crop_coeffs[i], 4, soilPars, _rooting_depth[i], _lai_factors[i])
    data_set_all[i,0] = compute_root_zone_moisture(obtain_soil_moisture(current_head), _rooting_depth[i])
    print(data_set_all[i,0])
    data_set_all[i, 1]  = _irrig_rates[i]
    data_set_all[i, 2] = _crop_coeffs[i]
    data_set_all[i, 3]  = _ref_evaps[i]
    data_set_all[i, 4]  = _rooting_depth[i]
    data_set_all[i, 5]  = compute_root_zone_moisture(obtain_soil_moisture(next_head),_rooting_depth[i])
    print(data_set_all[i,5])
    current_head = next_head
    pass 
plt.plot(data_set_all[:,0])
data_set_all_lai = np.zeros((len(_irrig_rates), 7))
data_set_all_lai[:,0] = data_set_all[:,0]
data_set_all_lai[:,1] = data_set_all[:,1]
data_set_all_lai[:,2] = data_set_all[:,2]
data_set_all_lai[:,3] = data_set_all[:,3]
data_set_all_lai[:,4] = data_set_all[:,4]
data_set_all_lai[:,5] = _lai_factors[0:len(_irrig_rates)]
data_set_all_lai[:,6] = data_set_all[:,5]
np.savetxt("./open_loop_data/training_data_mz1.csv", data_set_all_lai, delimiter="," )
